// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/common/base/bloom_filter.h>
#include <pollux/testing/gtest_utils.h>
#include <pollux/testing/exec/util/assert_query_builder.h>
#include <pollux/plan/plan_builder.h>
#include <pollux/functions/lib/aggregates/tests/utils/AggregationTestBase.h>
#include <pollux/functions/sparksql/aggregates/register.h>
using namespace kumo::pollux::plan;

namespace kumo::pollux::functions::aggregate::sparksql::test {
    namespace {
        class BloomFilterAggAggregateTest
                : public aggregate::test::AggregationTestBase {
        public:
            void SetUp() override {
                AggregationTestBase::SetUp();
                registerAggregateFunctions("");
            }

            VectorPtr getSerializedBloomFilter(int32_t capacity) {
                BloomFilter bloomFilter;
                bloomFilter.reset(capacity);
                for (auto i = 0; i < 9; ++i) {
                    bloomFilter.insert(melon::hasher<int64_t>()(i));
                }
                std::string data;
                data.resize(bloomFilter.serializedSize());
                bloomFilter.serialize(data.data());
                return make_constant(StringView(data), 1, VARBINARY());
            }
        };
    } // namespace

    TEST_F(BloomFilterAggAggregateTest, basic) {
        auto vectors = {
            make_row_vector({
                make_flat_vector<int64_t>(
                    100, [](vector_size_t row) { return row % 9; })
            })
        };
        auto expected = {make_row_vector({getSerializedBloomFilter(4)})};
        testAggregations(vectors, {}, {"bloom_filter_agg(c0, 5, 64)"}, expected);
    }

    TEST_F(BloomFilterAggAggregateTest, bloomFilterAggArgument) {
        auto vectors = {
            make_row_vector({
                make_flat_vector<int64_t>(
                    100, [](vector_size_t row) { return row % 9; })
            })
        };

        auto expected1 = {make_row_vector({getSerializedBloomFilter(3)})};
        testAggregations(vectors, {}, {"bloom_filter_agg(c0, 6)"}, expected1);

        // This capacity is kMaxNumBits / 16.
        auto expected2 = {make_row_vector({getSerializedBloomFilter(262144)})};
        testAggregations(vectors, {}, {"bloom_filter_agg(c0)"}, expected2);
    }

    TEST_F(BloomFilterAggAggregateTest, emptyInput) {
        auto vectors = {make_row_vector({make_flat_vector<int64_t>({})})};
        auto expected = {make_row_vector({make_null_constant(TypeKind::VARBINARY, 1)})};
        testAggregations(vectors, {}, {"bloom_filter_agg(c0, 5, 64)"}, expected);
    }

    TEST_F(BloomFilterAggAggregateTest, nullBloomFilter) {
        auto vectors = {make_row_vector({make_all_null_flat_vector<int64_t>(2)})};
        auto expectedFake = {
            make_row_vector(
                {make_nullable_flat_vector<StringView>({std::nullopt}, VARBINARY())})
        };
        POLLUX_ASSERT_THROW(
            testAggregations(
                vectors, {}, {"bloom_filter_agg(c0, 5, 64)"}, expectedFake),
            "First argument of bloom_filter_agg cannot be null");
    }

    TEST_F(BloomFilterAggAggregateTest, config) {
        auto vector = {
            make_row_vector({
                make_flat_vector<int64_t>(
                    100, [](vector_size_t row) { return row % 9; })
            })
        };
        std::vector<RowVectorPtr> expected = {
            make_row_vector({getSerializedBloomFilter(100)})
        };

        // This config will decide the bloom filter capacity, the expected value is
        // the serialized bloom filter, it should be consistent.
        testAggregations(
            vector,
            {},
            {"bloom_filter_agg(c0)"},
            expected,
            {{core::QueryConfig::kSparkBloomFilterMaxNumBits, "1600"}});

        // Test fails without setting the config.
        auto planNode = PlanBuilder(pool())
                .values(vector)
                .partialAggregation({}, {"bloom_filter_agg(c0)"})
                .finalAggregation()
                .planNode();
        auto actual = exec::test::AssertQueryBuilder(planNode).copyResults(pool());
        EXPECT_FALSE(
            expected[0]->childAt(0)->equal_value_at(actual->childAt(0).get(), 0, 0));
    }
} // namespace kumo::pollux::functions::aggregate::sparksql::test
